package com.snda.lidejia.mongodbtest;

import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.Mongo;
import java.net.UnknownHostException;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Hello world!
 *
 */
public class Insert {

    private static int THREAD_NUM = 20;
    private static AtomicLong id = new AtomicLong(0);
    private static final String HOST = "10.1.9.97";
    //
    private static final String buffer;

    static {
        StringBuilder sb = new StringBuilder();
        Random random = new Random();
        for (int i = 0; i < 10 * 1024; i++) {
            sb.append(random.nextInt(10));
        }
        buffer = sb.toString();
    }

    public static void main(String[] args) throws UnknownHostException, InterruptedException {
        if (args.length == 2) {
            THREAD_NUM = Integer.parseInt(args[0]);
            id.set(Long.parseLong(args[1]));
        }
        System.out.println(String.format("Thread number is %s", THREAD_NUM));
        System.out.println(String.format("Start index is %s", id.get()));
        //
        System.setProperty("MONGO.POOLSIZE", Integer.toString(THREAD_NUM));
        //
        ExecutorService es = Executors.newCachedThreadPool();
        final CountDownLatch activeThreadCounter = new CountDownLatch(THREAD_NUM);
        final Mongo mongo = new Mongo(HOST);
        try {
            for (int i = 0; i < THREAD_NUM; i++) {
                es.execute(new Runnable() {

                    public void run() {
                        try {
                            DB db = mongo.getDB("test");
                            DBCollection collection = db.getCollection("test");
                            while (true) {
                                BasicDBObject doc = new BasicDBObject();
                                doc.put("_id", id.getAndIncrement());
                                doc.put("name", buffer);
                                collection.insert(doc);
                            }
                        } finally {
                            activeThreadCounter.countDown();
                        }
                    }
                });
            }
            es.shutdown();
            activeThreadCounter.await();
        } finally {
            mongo.close();
        }
    }
}
